from lxml import etree
from tqdm import tqdm as tqdm
import os

# D:\datasets\YQSL-train-v3\crops\222\1
xml_dir = r'/Users/edy/Data/RC/label_data/xianchang_test/2022-02-11/test/Annotations'
crops_dir = r'/Users/edy/Data/RC/label_data/xianchang_test/2022-02-11/test/crops/rec/todo'
objname ='K1AYQSL00450001000 元气森林-元气森林-王林青苹果-苏打气泡水-PET瓶-480ml'


def recursive_parse_xml_to_dict(xml):
    if not xml:
        return {xml.tag: xml.text}
    result = {}
    for child in xml:
        child_result = recursive_parse_xml_to_dict(child)
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


from xml.dom.minidom import Document


class VOC_XML():
    def __init__(self, image_path, image_width, image_height):
        self.image_path = image_path
        self.image_width = image_width
        self.image_height = image_height
        self.init_base_info()

    def init_base_info(self):
        self.doc = Document()  # 创建DOM文档对象
        self.root_element = self.doc.createElement('annotation')  # 创建根元素
        self.root_element.setAttribute('verified', 'no')
        self.doc.appendChild(self.root_element)
        # 创建folder节点
        folder = self.doc.createElement('folder')
        folder_text = self.doc.createTextNode('JPEGImages')
        folder.appendChild(folder_text)
        self.root_element.appendChild(folder)
        # 创建filename节点
        filename = self.doc.createElement('filename')
        filename_text = self.doc.createTextNode('000001')
        filename.appendChild(filename_text)
        self.root_element.appendChild(filename)
        # 创建path节点
        path = self.doc.createElement('path')
        path_text = self.doc.createTextNode(self.image_path)
        path.appendChild(path_text)
        self.root_element.appendChild(path)
        # 创建source节点
        source = self.doc.createElement('source')
        # 创建source节点的子节点database
        database = self.doc.createElement('database')
        database_text = self.doc.createTextNode('Unknown')
        database.appendChild(database_text)
        source.appendChild(database)
        self.root_element.appendChild(source)
        # 创建图片的size节点
        size = self.doc.createElement('size')
        # 创建size节点的子节点width节点
        width = self.doc.createElement('width')
        width_text = self.doc.createTextNode(str(int(self.image_width)))
        width.appendChild(width_text)
        size.appendChild(width)
        # 创建size节点的子节点width节点
        height = self.doc.createElement('height')
        height_text = self.doc.createTextNode(str(int(self.image_height)))
        height.appendChild(height_text)
        size.appendChild(height)
        # 创建size节点的子节点depth节点
        depth = self.doc.createElement('depth')
        depth_text = self.doc.createTextNode("3")
        depth.appendChild(depth_text)
        size.appendChild(depth)
        self.root_element.appendChild(size)
        # 创建segmented节点
        segmented = self.doc.createElement('segmented')
        segmented_text = self.doc.createTextNode('0')
        segmented.appendChild(segmented_text)
        self.root_element.appendChild(segmented)

    def add_object_element(self, cls, box, x_axis_first=True, _score=None):
        """
        :param cls: 类别名
        :param box: 坐标值
        :param x_axis_first:  True: [xmin,ymin,xmax,ymax]    False:[ymin,xmin,ymax,xmax]
        :param _score: 是否有，主要自动标注时可能出现
        :return:
        """
        if x_axis_first:
            _xmin, _ymin, _xmax, _ymax, = box
        else:
            _ymin, _xmin, _ymax, _xmax = box
        object_element = self.doc.createElement('object')
        # 创建object节点的子节点name节点
        name = self.doc.createElement('name')
        name_text = self.doc.createTextNode(cls)
        name.appendChild(name_text)
        object_element.appendChild(name)
        if not _score is None:
            # 创建object节点的子节点score节点
            score = self.doc.createElement('score')
            score_text = self.doc.createTextNode(str(_score))
            score.appendChild(score_text)
            object_element.appendChild(score)
        # 创建object节点的子节点pose节点
        pose = self.doc.createElement('pose')
        pose_text = self.doc.createTextNode('Unspecified')
        pose.appendChild(pose_text)
        object_element.appendChild(pose)
        # 创建object节点的子节点truncated节点
        truncated = self.doc.createElement('truncated')
        truncated_text = self.doc.createTextNode('0')
        truncated.appendChild(truncated_text)
        object_element.appendChild(truncated)
        # 创建object节点的子节点Difficult节点
        Difficult = self.doc.createElement('difficult')
        Difficult_text = self.doc.createTextNode('0')
        Difficult.appendChild(Difficult_text)
        object_element.appendChild(Difficult)
        # 创建object节点的子节点bndbox节点
        bndbox_element = self.doc.createElement('bndbox')
        # 创建bndbox的子节点xmin

        xmin = self.doc.createElement('xmin')
        xmin_text = self.doc.createTextNode(str(int(float(_xmin))))
        xmin.appendChild(xmin_text)
        bndbox_element.appendChild(xmin)
        # 创建bndbox的子节点ymin
        ymin = self.doc.createElement('ymin')
        ymin_text = self.doc.createTextNode(str(int(float(_ymin))))
        ymin.appendChild(ymin_text)
        bndbox_element.appendChild(ymin)
        # 创建bndbox的子节点xmax
        xmax = self.doc.createElement('xmax')
        xmax_text = self.doc.createTextNode(str(int(float(_xmax))))
        xmax.appendChild(xmax_text)
        bndbox_element.appendChild(xmax)
        # 创建bndbox的子节点ymax
        ymax = self.doc.createElement('ymax')
        ymax_text = self.doc.createTextNode(str(int(float(_ymax))))
        ymax.appendChild(ymax_text)
        bndbox_element.appendChild(ymax)
        bndbox_element.appendChild(ymax)
        object_element.appendChild(bndbox_element)
        self.root_element.appendChild(object_element)

    def write_file(self, xml_path):
        with open(xml_path, 'w', encoding='utf-8') as f:
            f.write(self.doc.toprettyxml(indent='  '))


if __name__ == '__main__':
    for i in tqdm(
            filter(lambda x: x.endswith('.jpg'), os.listdir(crops_dir))):  # 根据文件夹里crop出来图片的命名，找到对应标签文件，对类别进行修改，直接覆盖原文件
        # print(i)
        if "frame" not in i:
            name=i.split('_')[0] #+'_'+i.split('_')[1] #+'_'+i.split('_')[2]  #原来的
        else:
            name = i.split('_')[0] + "_" + i.split('_')[1]  # +'_'+i.split('_')[1] #+'_'+i.split('_')[2]    #适应DP的

        y1, x1, y2, x2, hz = i.split('_')[-5:]
        xml_path = os.path.join(xml_dir, name + '.xml')
        print(xml_path)
        with open(xml_path, 'rb') as f:
            xml_str = f.read()

        xml = etree.fromstring(xml_str)
        xml_dict = recursive_parse_xml_to_dict(xml)
        if 'path' not in xml_dict["annotation"]:
            xml_dict["annotation"]['path'] = xml_dict["annotation"]['filename']
        xml_object = VOC_XML(xml_dict["annotation"]["path"],
                             xml_dict["annotation"]["size"]["width"],
                             xml_dict["annotation"]["size"]["height"])
        if 'object' in xml_dict["annotation"]:
            for index, obj in enumerate(xml_dict["annotation"]["object"]):

                if [obj['bndbox']['xmin'], obj['bndbox']['ymin'], obj['bndbox']['xmax'], obj['bndbox']['ymax']] == [x1,
                                                                                                                    y1,
                                                                                                                    x2,
                                                                                                                    y2]:
                    obj['name'] = objname
                    # continue 
                xml_object.add_object_element(obj['name'],
                                              [obj['bndbox']['xmin'], obj['bndbox']['ymin'],
                                               obj['bndbox']['xmax'],
                                               obj['bndbox']['ymax']])

        # xml_object.write_file(os.path.join('out_xml', xml_dir, os.path.basename(xml_path))) #raw,but seems no out_xml
        xml_object.write_file(os.path.join( xml_dir, os.path.basename(xml_path)))
